from major.models_manager import embedding_model,chat_model

from langchain_chroma import Chroma

# 直接实例化 Chroma 传入 向量db的路径即可

vector_db = Chroma(persist_directory="./chroma_db",
                   embedding_function=embedding_model.get_model() #
                   )
llm = chat_model.get_model()
def generate(query):
    # 2.检索
    docs = vector_db.similarity_search(query, k=3)

    # 3.生成
    # 先把pagecontent和metadata拼接起来
    docs = ['内容:'+doc.page_content + '\n' + '来源:'+doc.metadata['source'] for doc in docs]
    #print(docs)
    doc_context = '\n\n'.join(docs)
    #print(doc_context)
    # 拼接prompt
    prompt = f'''
    根据文档回答用户问题，如果文档中不包含用户问题的答案，请回答不知道。
    文档：
    {doc_context}
    用户问题：{query}
    '''
    #print(prompt)
    res = llm.stream(prompt)
    for i in res:
        print(i.content, end='')


generate("狗是什么")


